/*
参考github源码: https://github.com/NeowayLabs/es-reindex
*/
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/pborman/uuid"
	"io"
	"io/ioutil"
	"os"
	"time"

	"github.com/NeowayLabs/logger"
	"github.com/olivere/elastic/v7"
)

const defaultFromESHost = "http://127.0.0.1:9200"
const defaultToESHost = "http://127.0.0.1:9200"

var (
	fromHost       string
	fromIndex      string
	toHost         string
	toIndex        string
	newMapping     string
	mappingContent map[string]interface{}
	bulkSize       int
	forceYes       bool
	forceNo        bool

	esUserName string
	esPassword string

	useOriginDocId bool

	genOriginalMappingFile bool
)

// Notice 生成 newMapping 文件的方法

func main() {
	flag.StringVar(&fromHost, "fromHost", defaultFromESHost, "elastic search host to get data from")
	flag.StringVar(&fromIndex, "fromIndex", "", "name of index to reindex/copy")
	flag.StringVar(&toHost, "toHost", defaultToESHost, "elastic search host to get data from")
	flag.StringVar(&toIndex, "toIndex", "", "name of new-index")
	flag.StringVar(&newMapping, "newMapping", "", "path to new mapping file of new-index")
	flag.IntVar(&bulkSize, "bulkSize", 500, "amount of data to get in each request")
	flag.BoolVar(&forceYes, "forceYes", false, "even if destination alias already exists continue reindexing")
	flag.BoolVar(&forceNo, "forceNo", false, "if destination alias already exists abort reindexing")
	flag.StringVar(&esUserName, "esUserName", "", "username for ES Basic Auth")
	flag.StringVar(&esPassword, "esPassword", "", "password for ES Basic Auth")
	flag.BoolVar(&useOriginDocId, "useOriginDocId", true, "whether to use the original document ID")
	flag.BoolVar(&genOriginalMappingFile, "genOriginalMappingFile", false, "whether to genOriginalMappingFile")

	flag.Parse()

	ctx := context.Background()

	if fromIndex == "" {
		logger.Error("The `-fromIndex` parameters are required to reindex/copy")

		flag.Usage()
		os.Exit(1)
	}

	// Read new mapping file
	if newMapping != "" {
		mappingBytes, err := ioutil.ReadFile(newMapping)
		if err != nil {
			logger.Fatal("Error reading mapping file: %+v", err.Error())
		}
		// JSON 解析
		var mappingMap map[string]interface{}
		err = json.Unmarshal(mappingBytes, &mappingMap)
		if err != nil {
			logger.Fatal("Error parsing mapping file: %+v", err.Error())
		}
		// 清理非法字段
		if settings, ok := mappingMap["settings"].(map[string]interface{}); ok {
			if index, ok := settings["index"].(map[string]interface{}); ok {
				delete(index, "provided_name")
				delete(index, "version")
				delete(index, "creation_date")
				delete(index, "uuid")
			}
		}
		// 将清理后的内容传入
		mappingContent = mappingMap
		logger.Info("mappingContent: ", mappingMap)
		logger.Info("使用新 mapping 文件 %s", newMapping)
	}

	// Connect to clients
	fromClient, err := getESClient(fromHost)
	if err != nil {
		logger.Fatal("Error connecting to `%s`: %+v", fromHost, err.Error())
	}

	// 仅生成mapping文件
	if genOriginalMappingFile == true {
		GenOriginalMappingFile(ctx, fromClient)
	}

	toClient, err := getESClient(toHost)
	if err != nil {
		logger.Fatal("Error connecting to `%s`: %+v", toHost, err.Error())
	}

	// Set default toIndex
	if toIndex == "" {
		toIndex = fromIndex + "-" + uuid.New()[24:]
	}

	// Verify if fromIndex exists
	exists, err := fromClient.IndexExists(fromIndex).Do(ctx)
	if err != nil {
		logger.Fatal("Error verifying if index <%s> exists: %+v", fromIndex, err.Error())
	}
	if !exists {
		logger.Fatal("The index <%s> doesn't exists, we need a valid index or alias", fromIndex)
	}

	// Verify if toIndex already exists
	exists, err = toClient.IndexExists(toIndex).Do(ctx)
	if err != nil {
		logger.Fatal("Error verifying if index <%s> exists: %+v", toIndex, err.Error())
	}

	// Notice If toIndex don't exists we need create it
	if !exists {
		indexService := toClient.CreateIndex(toIndex)
		// If -newMapping was not provided use original mapping
		if newMapping == "" {
			mapping, err := fromClient.GetMapping().Index(fromIndex).Do(ctx)
			if err != nil {
				logger.Fatal("Error getting mapping of index <%s>", fromIndex)
			}

			for _, v := range mapping {
				mapping = v.(map[string]interface{})
				break
			}

			resp, err := fromClient.IndexGetSettings().Index(fromIndex).Do(ctx)
			if err != nil {
				logger.Fatal("Error getting settings of index <%s>", fromIndex)
			}

			for _, v := range resp {
				mapping["settings"] = v.Settings
				break
			}

			// 清理 settings.index 下非法字段
			if settings, ok := mapping["settings"].(map[string]interface{}); ok {
				if indexSettings, ok := settings["index"].(map[string]interface{}); ok {
					delete(indexSettings, "creation_date")
					delete(indexSettings, "uuid")
					delete(indexSettings, "version")
					delete(indexSettings, "provided_name")
				}
			}

			logger.Info("使用原mapping结构........")
			indexService.BodyJson(mapping)
		} else {
			logger.Info("使用文件..............")
			indexService.BodyJson(mappingContent)
		}

		// Notice 如果目标索引已存在，且未设置 -force-yes，则提示用户输入是否继续。
		createNewIndex, err := indexService.Do(ctx)
		if err != nil {
			logger.Fatal("Error creating new index <%s>: %+v", toIndex, err.Error())
		}
		if !createNewIndex.Acknowledged {
			logger.Fatal("Was not possible create new index <%s>", toIndex)
		}

		logger.Info("New index <%s> was created!", toIndex)
	} else if forceNo {
		logger.Fatal("Index <%s> already exists in destination server", toIndex)
	} else if !forceYes {
		if !askForConfirmation(fmt.Sprintf("Index <%s> already exists, do you want index all documents without change the current mapping? (yes/no) ", toIndex)) {
			os.Exit(0)
		}
	}

	// Notice Reindex fromIndex to toIndex
	// 定义 scroll 查询
	scroll := fromClient.Scroll(fromIndex).
		Size(bulkSize).
		FetchSource(true)

	logger.Info("开始滚动读取文档进行 Reindex...")

	var totalDocs int
	reindexStart = time.Now()

	for {
		results, err := scroll.Do(ctx)
		// 所有文档读取完毕
		if err == io.EOF {
			logger.Info("所有文档读取完毕......")
			break
		}
		if err != nil {
			logger.Fatal("scroll 查询失败: %+v", err)
		}

		if len(results.Hits.Hits) == 0 {
			break
		}

		bulkRequest := toClient.Bulk()
		for _, hit := range results.Hits.Hits {
			var doc map[string]interface{}
			if err := json.Unmarshal(hit.Source, &doc); err != nil {
				logger.Error("无法解析文档 ID %s 的 source: %v", hit.Id, err)
				continue
			}
			// 是否与原始文档 ID 保持一致
			var req elastic.BulkableRequest
			if useOriginDocId {
				req = elastic.NewBulkIndexRequest().Index(toIndex).Id(hit.Id).Doc(doc)
			} else {
				req = elastic.NewBulkIndexRequest().Index(toIndex).Doc(doc)
			}
			bulkRequest = bulkRequest.Add(req)
		}

		bulkResp, err := bulkRequest.Do(ctx)
		if err != nil {
			logger.Fatal("bulk 写入失败: %+v", err)
		}

		if bulkResp.Errors {
			for _, item := range bulkResp.Failed() {
				logger.Error("写入失败：Index[%s] Id[%s]: %s", item.Index, item.Id, item.Error.Reason)
			}
		}

		totalDocs += len(results.Hits.Hits)
		logger.Info("已成功写入 %d 条文档... [耗时: %s]", totalDocs, time.Since(reindexStart))
	}
	logger.Info("Reindex 完成！共写入 %d 条文档，总耗时 %s", totalDocs, time.Since(reindexStart))

	// Notice If index is a alias, update its reference
	// Notice 如果原始索引 fromIndex 是一个别名（alias），那么程序会把这个别名重新指向你新建的目标索引 toIndex，而不再指向原索引，实现“别名切换”。
	aliasesService := toClient.Aliases()
	aliases, err := aliasesService.Do(ctx)
	if err != nil {
		logger.Fatal("Error getting aliases: %+v", err.Error())
	}
	indices := aliases.IndicesByAlias(fromIndex)
	fmt.Println("aliases.indices:>>>>>> ", len(indices), indices)
	if len(indices) > 0 {
		aliasService := toClient.Alias()
		for _, index := range indices {
			aliasService.Remove(index, fromIndex)
		}
		_, err = aliasService.Add(toIndex, fromIndex).Do(ctx)
		if err != nil {
			logger.Fatal("Error updating alias <%s>:  %+v", fromIndex, err.Error())
		}

		logger.Info("Alias <%s>: %+v was removed and now point to: <%s>", fromIndex, indices, toIndex)
	}

}

func getESClient(esURL string) (*elastic.Client, error) {

	options := []elastic.ClientOptionFunc{
		elastic.SetURL(esURL),
		// 关闭sniff(客户端自动发现并更新集群中节点的列表，从而实现更智能的负载均衡和容错能力)
		elastic.SetSniff(false),
		elastic.SetErrorLog(logger.DefaultLogger.Handlers[0].(*logger.StdOutHandler).ErrorLogger),
		elastic.SetInfoLog(logger.DefaultLogger.Handlers[0].(*logger.StdOutHandler).DebugLogger),
		elastic.SetTraceLog(logger.DefaultLogger.Handlers[0].(*logger.StdOutHandler).DebugLogger),
	}
	// ES用户名与密码
	if esUserName != "" && esPassword != "" {
		options = append(options, elastic.SetBasicAuth(esUserName, esPassword))
	}

	esClient, err := elastic.NewClient(options...)
	if err != nil {
		return esClient, err
	}

	esVersion, err := esClient.ElasticsearchVersion(esURL)
	if err != nil {
		logger.Fatal("Error getting ES version: %+v", err.Error())
	}
	logger.Info("Connected in Elasticsearch <%s>, version %s", esURL, esVersion)

	return esClient, err
}

var (
	reindexProgress = -1
	reindexStart    time.Time
)

func showReindexProgress(current, total int64) {
	percent := (float64(current) / float64(total)) * 100
	if int(percent) > reindexProgress {
		reindexProgress = int(percent)
		logger.Info("Reindexing... %d%% [Time elapsed: %s]", reindexProgress, time.Since(reindexStart).String())
	}
}

// Based in: https://gist.github.com/albrow/5882501

// askForConfirmation uses Scanln to parse user input. A user must type in "yes" or "no" and
// then press enter. It has fuzzy matching, so "y", "Y", "yes", "YES", and "Yes" all count as
// confirmations. If the input is not recognized, it will ask again. The function does not return
// until it gets a valid response from the user. Typically, you should use fmt to print out a question
// before calling askForConfirmation. E.g. fmt.Println("WARNING: Are you sure? (yes/no)")
func askForConfirmation(message string) bool {
	fmt.Printf(message)

	var response string
	_, err := fmt.Scanln(&response)
	if err != nil {
		logger.Fatal("Cannot read from stdin: %s", err.Error())
	}

	yesResponses := []string{"y", "Y", "yes", "Yes", "YES"}
	noResponses := []string{"n", "N", "no", "No", "NO"}

	if containsString(yesResponses, response) {
		return true
	} else if containsString(noResponses, response) {
		return false
	} else {
		return askForConfirmation(message)
	}
}

// You might want to put the following two functions in a separate utility package.

// posString returns the first index of element in slice.
// If slice does not contain element, returns -1.
func posString(slice []string, element string) int {
	for index, elem := range slice {
		if elem == element {
			return index
		}
	}
	return -1
}

// containsString returns true iff slice contains element
func containsString(slice []string, element string) bool {
	return !(posString(slice, element) == -1)
}

func GenOriginalMappingFile(ctx context.Context, fromClient *elastic.Client) {

	mapping, err := fromClient.GetMapping().Index(fromIndex).Do(ctx)
	if err != nil {
		logger.Fatal("Error getting mapping of index <%s>", fromIndex)
	}

	for _, v := range mapping {
		mapping = v.(map[string]interface{})
		break
	}

	resp, err := fromClient.IndexGetSettings().Index(fromIndex).Do(ctx)
	if err != nil {
		logger.Fatal("Error getting settings of index <%s>", fromIndex)
	}

	for _, v := range resp {
		mapping["settings"] = v.Settings
		break
	}

	// 清理 settings.index 下非法字段
	if settings, ok := mapping["settings"].(map[string]interface{}); ok {
		if indexSettings, ok := settings["index"].(map[string]interface{}); ok {
			delete(indexSettings, "creation_date")
			delete(indexSettings, "uuid")
			delete(indexSettings, "version")
			delete(indexSettings, "provided_name")
		}
	}

	// Notice 将原index的mapping结构写入到一个json文件
	mappingBytes, err := json.MarshalIndent(mapping, "", "  ")
	if err != nil {
		logger.Fatal("Error marshalling mapping to JSON: %+v", err)
	}
	err = ioutil.WriteFile("original_mapping.json", mappingBytes, 0644)
	if err != nil {
		logger.Fatal("Error writing mapping JSON file: %+v", err)
	}
	logger.Info("使用原mapping结构并已写入 original_mapping_org.json 文件")

	// Notice 直接停止进程
	os.Exit(1)
}
